# # 
#  Session.py
#  ExpectIT
#  
#  Author: Marvin S. Addison
#  Date: 2007-10-15
#  Version: $Revision: $
#  Updated: $Date: $
#  Copyright 2008 Marvin S. Addison. All rights reserved.
#  See LICENSE.txt for licensing information.
#

import pexpect
import util
import Auth
import Command
import re
import getpass
from util import printerr, quit

"""Maximum number of allowed interactive auth attempts."""
MAX_AUTH_TRIES = 3

"""Password prompt string."""
PASSWORD_PROMPT = "Password:"

class Session:
	"""Session objects represent interactive sessions on remote hosts."""

	def __init__(self):
		"""Initializes a new Session object."""
		self.__exp = None
		self.__expect_re = None
		self.__prompt_re = None
		self.host = ''
		self.protocol = 'ssh'
		self.connect_timeout = 10
		self.command_timeout = 3
		self.auth = None
		self.commands = []
		self.prompt = ''
		self.echo = False
		self.separator = '\n\n\n'
		self.is_connected = False
		self.captures = {}

	def connect(self):
		"""Initiates an interactive session with the remote host."""
		cmd = '%s -l%s' % (self.protocol, self.auth.user)
		if self.auth.method == 'key':
			cmd += ' -i%s' % self.auth.credential
		cmd += ' %s' % self.host
		printerr('Connecting with command "%s"' % cmd)
		self.__exp = pexpect.spawn(cmd, timeout=self.connect_timeout)
		if self.auth.method == 'password':
			try:
				self.__exp.expect(PASSWORD_PROMPT, timeout=self.command_timeout)
				printerr('Got password prompt. Sending password.')
				self.__exp.sendline(self.auth.credential)
			except pexpect.TIMEOUT:
				quit('Never received password prompt.')
		n_tries = 0
		while (n_tries <= MAX_AUTH_TRIES and not self.is_connected):
			n_tries += 1
			try:
				self.__exp.expect(self.prompt, timeout=self.command_timeout)
				self.is_connected = True
				printerr('Connected.')
			except pexpect.TIMEOUT:
				printerr('Authentication failure.')
				if n_tries > MAX_AUTH_TRIES:
					quit('Giving up after %s tries.' % MAX_AUTH_TRIES)
				else:
					printerr('Trying interactive auth %s of %s:' % \
						(n_tries, MAX_AUTH_TRIES))
					self.__exp.sendline(getpass.getpass('Password: '))

	def execute(self):
		"""Executes all the commands on the remote host and returns the results
		of each on standard output.  Commands are separated by the separator
		string defined in the config file, if defined, otherwise 3 blank lines."""
		self.__prompt_re = re.compile(self.prompt)
		i = 0
		for cmd in self.commands:
			self.__expect_re = None
			if cmd.expect:
				self.__expect_re = re.compile(cmd.expect)
			for c in self.eval_captures(cmd):
				self.exec_cmd(c, i>0)
				i += 1


	def disconnect(self):
		"""Disconnects from the remote hosts and ends the session."""
		try:
			printerr('\nDisconnecting from %s' % self.host)
			self.__exp.sendline('\cC')
			self.__exp.sendline('exit')
		except pexpect.TIMEOUT:
			self.__exp.expect('.*', timeout=self.command_timeout)
		self.__exp.terminate()
		self.is_connected = False

	def store_captures(self, s):
		"""Store any captures from evaluating the given string against the
		regex, which may contain capture expressions."""
		match = self.__expect_re.search(s)
		if match:
			groups = match.groupdict()
			for key in groups:
				if key in self.captures:
					self.captures[key].append(groups[key])
				else:
					self.captures[key] = [groups[key]]

	def eval_captures(self, cmd):
		"""Replaces any occurances of <name> in the given command where
		name is a named capture expression.  Returns a list of commands
		as named capture expressions can contain multiple values."""
		commands = []
		for name in self.captures.keys():
			parm = '<%s>' % name;
			if cmd.command.find(parm) > -1:
				for value in self.captures[name]:
					clone = cmd.clone()
					clone.command = cmd.command.replace(parm, value)
					commands.append(clone)
		if len(commands) == 0:
			commands.append(cmd)
		return commands
	
	def exec_cmd(self, command, print_separator=True):
		"""Execute the given command."""	
		self.__exp.sendline(command.command)
		if print_separator:
			print self.separator
		printerr('Executing command: %s' % command)
		for line in readlines(self.__exp):
			if command.expect:
				self.store_captures(line)
			if not self.echo and line.startswith(command.command):
				pass
			elif self.__prompt_re.search(line):
				pass
			elif command.show_output:
				print line


def readlines(exp):
	"""Reads the output from a pexpect object that presumably has
	an internal buffer containing data resulting from the execution
	of some command.  Output is returned as an array of lines."""
	size = 1024
	output = ''
	try:
		while True:
			output += exp.read_nonblocking(size, 1)
	except pexpect.TIMEOUT:
		pass
	return output.split('\r\n')


def from_dict(dict):
	"""Creates a new Session object initialized with properties from the
	supplied dictionary."""
	session = Session()
	util.load_object(session, dict, exclude_keys=['auth', 'commands'])
	if dict.has_key('auth'):
		session.auth = Auth.from_dict(dict['auth'])
	for command in dict['commands']:
		session.commands.append(Command.from_dict(command))
	return session
